from astro.coord import match
from astro.utilities import append_field
from itertools import count,izip
from astro.utilities import read_ned_batch

class Selector():
    def __init__(self,ax,qsos):
        self.marker, = ax.plot([0],[100],'o',mfc=None,mec='g')
        self.ax = ax
        self.q = qsos
    def onpick(self,event):
        ax,q,marker = self.ax,self.q,self.marker
        i = event.ind
        ax.set_title('%s: z%s mag%s i%s' % (q.name[i],q.z[i],q.mag[i],i))
        marker.set_xdata(q.ra[i])
        marker.set_ydata(q.dec[i])
        pl.draw()

# read veron QSOs, all with v < 19.0
def read_veron():
    """Return Veron QSOs."""
    temp = open('veron_qsos/veron2006.tsv')
    f = temp.readlines()
    temp.close()
    f = [r for r in f if r]
    f = [r for r in f if len(r) > 1]
    print 'reading in Veron catalogue...'
    qsos = readtxt(f[41:], sep='|', names='recno,vclass,radio,name,'
                   'sra,sdec,aor,F6cm,r_F6cm,F20cm,r_F20cm,lz,z,'
                   'sclass,nvmag,mag,bmv,umb,Mabs,FC,r_Vmag,refz')
    print 'selecting only objects with redshifts and magnitudes...'
    temp = [list(q) + list(s2dec(q.sra,q.sdec)) for q in qsos]
    # select only objects with redshifts (some blacs don't have one) and mags
    temp = [r for r in temp if r[12] != '' and r[15] != '']
    temp = [r[:12]+[float(r[12])]+r[13:15]+[float(r[15])] + r[16:] for r in temp]
    print 'creating final record array...'
    qsos = np.rec.fromrecords(temp,names='recno,vclass,radio,name,sra,sdec,aor,'
                              'F6cm,r_F6cm,F20cm,r_F20cm,lz,z,'
                              'sclass,nvmag,mag,bmv,umb,Mabs,FC,r_Vmag,'
                              'refz,ra,dec')
    return qsos

# read simbad QSOs, all with v < 18.5
def read_simbad():
    """Return Simbad QSOs."""
    temp = open('simbad/simbad_qsos_vmag_lt_18.5.txt')
    f = temp.readlines()
    temp.close()
    f = f[9:-2]
    name,radecs,z,umag,Bmag,Vmag = readtxt(f[2:],usecols=(1,3,4,5,6,7),sep='|')
    ras,decs = zip(*[rd.split() for rd in radecs])
    ras = [float(r) for r in ras]
    decs = [float(d) for d in decs]
    sras,sdecs = zip(*[dec2s(r,d) for r,d in zip(ras,decs)] ) 
    qsos = np.rec.fromarrays([name,ras,decs,sras,sdecs,z,umag,Bmag,Vmag],
                             names='name,ra,dec,sra,sdec,z,umag,Bmag,mag')
    return qsos

def read_gto():
    """Return GTO targets (QSOs)."""
    temp = open('gto_qso_targets.txt')
    f = temp.readlines()
    temp.close()
    f = [r for r in f if r]
    f = [r.split() for r in f if r[0] != '#']
    cols = zip(*f)
    sra = [':'.join(c) for c in zip(*cols[2:5])]
    sdec = [':'.join(c) for c in zip(*cols[5:8])]
    ra,dec = zip(*[s2dec(r,d) for r,d in izip(sra,sdec)])
    orbits,name,vmag = ([int(o) for o in cols[0]],cols[1],
                        [float(m) for m in cols[10]])
    prog = [2045]*4 + [2019]*4 + [2034]*10 + [2035]*17 + [2044]*18
    gto = np.rec.fromarrays([name,ra,sra,dec,sdec,orbits,vmag,prog],
                            names='name,ra,sra,dec,sdec,orbits,mag,prog')
    return gto

def plotall(qsos):
    ras,decs = qsos.ra,qsos.dec    
    pl.ioff()
    pl.close()
    fig = pl.figure()
    a = fig.add_subplot(111)   
    ########################################################################
    #GDDS field centres. field size 5.5 x 5.5 arcmin
    ########################################################################
    gradecs = [('02:09:41.30',' -04:37:54.0'),
               ('12:05:22.17',' -07:22:27.9'),
               ('15:23:47.83',' -00:05:21.1'),
               ('22:17:41.0','+00:15:20.0')]
    gra,gdec = zip(*[s2dec(*rd) for rd in gradecs])
    temp = []
    # plot 5.5 by 5.5 arcmin square around each GDDS pointing
    for i,r,d in zip(count(0),gra,gdec):
        w = h =  5.5/60
        r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
        temp.append((r1<ras)&(ras<r2)&(d1<decs)&(decs<d2))
        a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],'k',lw=2)
        a.plot([r],[d],'kx',ms=10,mew=2,alpha=0.7)
    gcond = temp[0]
    for item in temp[1:]:
        gcond = gcond | item

    ########################################################################
    #zCOSMOS bright field (everything with I<22. field size 1.3 x 1.24 deg)
    ########################################################################
    r,d = s2dec('10 00 28','02 13 33')
    w,h = 1.3,1.24
    r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
    zcosmos_cond = (r1<ras)&(ras<r2)&(d1<decs)&(decs<d2)
    a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],'g',lw=2)

    #################################################################
    # VVDS fields. All 2x2 deg except 03 h field (CDFS) which is 2x1.
    #################################################################
    vradec = [('02:26:00','-04:30:00'), ('03:32:28','-27:48:30'),
              ('10:03:39','01:54:39 '), ('14:00:00','05:00:00 '),
              ('22:17:51','00:24:27 ')] 
    vra,vdec = zip(*[s2dec(*rd) for rd in vradec])
    r1,r2 = 36.260, 37.071               # Initial 1x1 deg VVDS DEEP field
    d1,d2 = -4.821, -4.139
    a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],'k',lw=2)

    temp = []
    # plot square around each VVDS pointing
    for i,r,d in zip(count(0),vra,vdec):
        w,h =  (2.,1.) if i == 1 else (2.,2.)
        r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
        temp.append((r1<ras)&(ras<r2)&(d1<decs)&(decs<d2))
        a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],'k',lw=2)
    vcond = temp[0]
    for item in temp[1:]:
        vcond = vcond | item

    #r,d = s2dec('3 32 39.0', '-27 47 29.1')    # UDF (centre)
    #a.plot([r],[d],'s',mec='r',mfc=None)

    ##################################################################
    # deep2 field centres. 2.0 x 0.27 deg, (14h field 2.0 x 0.15 deg)
    ##################################################################
    dradec = [('14 17 00','+52 30 00'),  # Groth Survey Strip region
              ('16 52 00','+34 55 00'),  # Zone of very low extinction
              ('23 30 00','+00 00 00'),  # On SDSS deep strip
              ('02 30 00','+00 00 00')]
    dra,ddec = zip(*[s2dec(*rd) for rd in dradec])
    # plot rectangle around each DEEP2 pointing
    temp = []
    for i,r,d in zip(count(0),dra,ddec):
        w,h =  (2.,0.15) if i == 0 else (2.,0.27)
        r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
        temp.append((r1<ras)&(ras<r2)&(d1<decs)&(decs<d2))
        a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],'k',lw=2)
    dcond = temp[0]
    for item in temp[1:]:
        dcond = dcond | item

    #########################################################
    # Stripe 82
    #########################################################
    d1,d2 = -1.266, 1.266
    r1,r2 = (20. + 34./60)*360./24, 4.*360./24
    # plot rectangles around stripe 82 region
    a.plot([0,r2,r2,0,0],[d1,d1,d2,d2,d1],'k',lw=2)
    a.plot([360,r1,r1,360,360],[d1,d1,d2,d2,d1],'k',lw=2)
    strcond = ((ras<r2)&(d1<decs)&(decs<d2)) | ((ras>r1)&(d1<decs)&(decs<d2))

    #########################################################
    # CFHT wide legacy survey
    #########################################################
    cradec = [('02 18 00','-07 00 00'), 
              ('08 54 00','-04 15 00'), 
              ('14 17 54','+54 30 31'),
              ('22 13 17.76','+01 19 00.0')]
    cra,cdec = zip(*[s2dec(*rd) for rd in cradec])
    
    # plot rectangles around CFHT wide region
    temp = []
    for i,r,d in zip(count(0),cra,cdec):
        if i == 0:
            w,h =  (9.,8.)
        elif i == 3:
            w,h =  (5.,5.)
        else:
            w,h = (7.,7.)
        r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
        temp.append((r1<ras)&(ras<r2)&(d1<decs)&(decs<d2))
        a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],color='0.3',lw=2)
    ccond = temp[0]
    for item in temp[1:]:
        ccond = ccond | item

    #########################################################
    # CFHT deep legacy survey
    #########################################################
    cradec = [('02h 26m 00s','-04d 30m 00s'), 
              ('10h 00m 29s','02d 12m 21s'),
              ('14h 19m 28s','+52d 40m 41s'),
              ('22h 15m 31s','-17d 44m 00s')]
    cra,cdec = zip(*[s2dec(*rd) for rd in cradec])
    
    # plot rectangles around CFHT deep region
    temp = []
    for i,r,d in zip(count(0),cra,cdec):
        w,h =  (1.,1.)
        r1,r2,d1,d2 = r-w/2,r+w/2,d-h/2,d+h/2
        temp.append((r1<ras)&(ras<r2)&(d1<decs)&(decs<d2))
        a.plot([r1,r2,r2,r1,r1],[d1,d1,d2,d2,d1],color='r',lw=2)
    cdcond = temp[0]
    for item in temp[1:]:
        cdcond = cdcond | item

    #########################################################
    # QSOs
    #########################################################

    # GTO QSOs (diamonds)
    a.plot(gto.ra,gto.dec,'D',mfc=None,mec='k')

    # Veron QSOs, colour by redshift
    q = qsos
    col = a.scatter(q.ra,q.dec,marker='o',s=10,
                    c=q.z,cmap=pl.cm.winter_r,faceted=False,
                    picker=True)

    a.set_xlabel('RA'); a.set_ylabel('Dec'); a.axis('equal')
    cbar = pl.colorbar(col)
    cbar.ax.set_ylabel('Redshift')
    # initialise the marker and selector routine
    selector = Selector(a,q)
    fig = pl.gcf()
    fig.canvas.mpl_connect('pick_event',selector.onpick)
    pl.show()
    return selector,dcond,vcond,zcosmos_cond,gcond,strcond,ccond,cdcond

# condition = ((ramin<qsos.ra)&(qsos.ra<ramax)&
#              (decmin<qsos.dec)&(qsos.dec<decmax))
# q_vvds_deep = qsos[condition]

qsos = read_veron()
#qsos = read_ned_batch('ned/ned_batchres_vvds.txt')
#qsos = read_simbad()
qsos = qsos[qsos.z>0.5]

gto = read_gto()
ind = match(gto.ra,gto.dec,qsos.ra,qsos.dec,8.0)
ver_info = [(list(qsos[i]) if i!=-1 else [-1]*16) for i in ind]
gto = np.rec.fromarrays(
    zip(*gto.tolist()) + zip(*ver_info), names='name,ra,sra,dec,sdec,orbits'
    ',vmag,prog,' + ','.join(['ver_' + n for n in qsos.dtype.names]))

# write out a list of GTO QSOs and header

# g = gto
# writecol(
#     'gto_qsos_tmp.txt',(g.name,g.sra,g.sdec,g.ver_z,g.vmag,g.ver_vmag,
#                         g.ver_class,g.ver_name,g.prog),sep='|',width='fixed')

# f = open('gto_qsos_tmp.header','w')
# f.write("""# GTO name
# # GTO ra (J2000)
# # GTO dec (J2000)
# # Veron redshift     (-1 means no match within 8 arcsec of Veron position)
# # GTO V mag
# # Veron mag        (-1 means no match)
# # Veron class (Q=QSO, B=BL Lac, A=AGN) (-1 means no match)
# # Veron name         (-1 means no match)
# # GTO program number
# """)
# f.close()

s,dcond,vcond,zc_cond,gcond,strcond,ccond,cdcond = plotall(qsos)

print 'Finished plotting'

# select the qsos within each survey.
mag_cond =  (qsos.mag>0)#& (qsos.mag<18)

vvds_qsos = qsos[vcond & mag_cond]
vvds_qsos = append_field(vvds_qsos,'survey',['VVDS']*len(vvds_qsos))

zcosmos_qsos = qsos[zc_cond & mag_cond]
zcosmos_qsos = append_field(zcosmos_qsos,'survey',
                            ['zCOSMOS']*len(zcosmos_qsos))
deep2_qsos = qsos[dcond & mag_cond]
deep2_qsos = append_field(deep2_qsos,'survey',['DEEP2']*len(deep2_qsos))
gdds_qsos = qsos[gcond & mag_cond]
gdds_qsos = append_field(gdds_qsos,'survey',['GDDS']*len(gdds_qsos))

s_qsos = qsos[strcond & mag_cond]
s_qsos = append_field(s_qsos,'survey', ['stripe82']*len(s_qsos))

c_qsos = qsos[ccond & mag_cond]
c_qsos = append_field(c_qsos,'survey', ['CFHTwide']*len(c_qsos))

cd_qsos = qsos[cdcond & mag_cond]
cd_qsos = append_field(cd_qsos,'survey', ['CFHTdeep']*len(cd_qsos))

# for i,q in enumerate([vvds_qsos,zcosmos_qsos,deep2_qsos,gdds_qsos]):
#     writecol('junk%s.dat' % i,(q.name,q.sra,q.sdec,q.z,q.mag,
#                                q.vclass,q.survey),sep='|',width='fixed')

# below for SIMBAD qsos list only

# for i,q in enumerate([vvds_qsos,zcosmos_qsos,deep2_qsos,gdds_qsos]):
#     writecol('sjunk%s.dat' % i,(q.name,q.sra,q.sdec,q.z,q.vmag,q.survey),
#              sep='|',width='fixed')

# f = open('survey_qsos_tmp.header','w')
# f.write("""# name
# # ra (J2000)
# # dec (J2000)
# # redshift
# # V mag
# # Veron class (Q=QSO, B=BL Lac, A=AGN)
# # Survey
# """)
# f.close()
